use std::collections::BTreeMap;
use std::path::Path;

use lopdf::{Document, Object, ObjectId, Bookmark};
use walkdir::WalkDir;


pub fn merge_pdfs_in_directory<P: AsRef<Path>>(directory: P, output_path: P) -> Result<(), Box<dyn std::error::Error>> {
    // 收集目录中的所有PDF文件
    let mut pdf_files: Vec<_> = WalkDir::new(directory)
        .into_iter()
        .filter_map(|e| e.ok())
        .filter(|e| {
            e.path()
                .extension()
                .map_or(false, |ext| ext.eq_ignore_ascii_case("pdf"))
        })
        .collect();

    // 按文件名排序
    pdf_files.sort_by_key(|a| a.path().to_path_buf());

    // 如果没有找到PDF文件，返回错误
    if pdf_files.is_empty() {
        return Err(Box::from(anyhow::anyhow!("No PDF files found in the directory")));
    }

    // 加载所有PDF文档
    let documents: Vec<Document> = pdf_files
        .iter()
        .filter_map(|entry| Document::load(entry.path()).ok())
        .collect();

    if documents.is_empty() {
        return Err(Box::from(anyhow::anyhow!("No valid PDF files found in the directory")));
    }

    // 开始合并过程
    let mut max_id = 1;
    let mut pagenum = 1;
    let mut documents_pages = BTreeMap::new();
    let mut documents_objects = BTreeMap::new();
    let mut document = Document::with_version("1.5");

    for mut doc in documents {
        let mut first = false;
        doc.renumber_objects_with(max_id);

        max_id = doc.max_id + 1;

        documents_pages.extend(
            doc.get_pages()
                .into_iter()
                .map(|(_, object_id)| {
                    if !first {
                        let bookmark = Bookmark::new(
                            format!("Page_{}", pagenum),
                            [0.0, 0.0, 1.0],
                            0,
                            object_id,
                        );
                        document.add_bookmark(bookmark, None);
                        first = true;
                        pagenum += 1;
                    }

                    (object_id, doc.get_object(object_id).unwrap().to_owned())
                })
                .collect::<BTreeMap<ObjectId, Object>>(),
        );
        documents_objects.extend(doc.objects);
    }

    // "Catalog" 和 "Pages" 是必需的
    let mut catalog_object: Option<(ObjectId, Object)> = None;
    let mut pages_object: Option<(ObjectId, Object)> = None;

    // 处理除"Page"类型外的所有对象
    for (object_id, object) in documents_objects.iter() {
        match object.type_name().unwrap_or(b"") {
            b"Catalog" => {
                catalog_object = Some((
                    if let Some((id, _)) = catalog_object {
                        id
                    } else {
                        *object_id
                    },
                    object.clone(),
                ));
            }
            b"Pages" => {
                if let Ok(dictionary) = object.as_dict() {
                    let mut dictionary = dictionary.clone();
                    if let Some((_, ref object)) = pages_object {
                        if let Ok(old_dictionary) = object.as_dict() {
                            dictionary.extend(old_dictionary);
                        }
                    }

                    pages_object = Some((
                        if let Some((id, _)) = pages_object {
                            id
                        } else {
                            *object_id
                        },
                        Object::Dictionary(dictionary),
                    ));
                }
            }
            b"Page" | b"Outlines" | b"Outline" => {} // 忽略这些类型
            _ => {
                document.objects.insert(*object_id, object.clone());
            }
        }
    }

    // 检查必需对象
    let pages_object = pages_object.ok_or_else(|| anyhow::anyhow!("Pages root not found"))?;
    let catalog_object = catalog_object.ok_or_else(|| anyhow::anyhow!("Catalog root not found"))?;

    // 处理所有页面对象
    for (object_id, object) in documents_pages.iter() {
        if let Ok(dictionary) = object.as_dict() {
            let mut dictionary = dictionary.clone();
            dictionary.set("Parent", pages_object.0);
            document.objects.insert(*object_id, Object::Dictionary(dictionary));
        }
    }

    // 更新Pages对象
    if let Ok(dictionary) = pages_object.1.as_dict() {
        let mut dictionary = dictionary.clone();
        dictionary.set("Count", documents_pages.len() as u32);
        dictionary.set(
            "Kids",
            documents_pages
                .into_iter()
                .map(|(object_id, _)| Object::Reference(object_id))
                .collect::<Vec<_>>(),
        );
        document.objects.insert(pages_object.0, Object::Dictionary(dictionary));
    }

    // 更新Catalog对象
    if let Ok(dictionary) = catalog_object.1.as_dict() {
        let mut dictionary = dictionary.clone();
        dictionary.set("Pages", pages_object.0);
        dictionary.remove(b"Outlines");
        document.objects.insert(catalog_object.0, Object::Dictionary(dictionary));
    }

    document.trailer.set("Root", catalog_object.0);

    // 更新文档
    document.max_id = document.objects.len() as u32;
    document.renumber_objects();
    document.adjust_zero_pages();
    
    if let Some(n) = document.build_outline() {
        if let Ok(Object::Dictionary(dict)) = document.get_object_mut(catalog_object.0) {
            dict.set("Outlines", Object::Reference(n));
        }
    }

    document.compress();

    // 保存合并后的文档
    document.save(output_path)?;

    Ok(())
}